#include "primitive.h"
#include "bsdf.h"

namespace csupbr {

    ///////////////////////////
    /// LambertianReflection
    ///////////////////////////
    
    Spectrum LambertianReflection::sample(BSDFRecord &bsdfRec, const Point2f &sample, Float &pdf) const {
        bsdfRec.wi = bsdfRec.isct.shCoord.toWorld(cosineHemisphereSampling(sample));
        bsdfRec.wi.normalize();
        pdf = this->pdf(bsdfRec);
        return {InvPi};
    }

    Spectrum LambertianReflection::evaluate(BSDFRecord &bsdfRec, Float &pdf) const {
        pdf = this->pdf(bsdfRec);
        return {InvPi};
    }

    Float LambertianReflection::pdf(const BSDFRecord &bsdfRec) const {
        return InvPi;
    }

    ///////////////////////////
    /// SpecularReflection
    ///////////////////////////

    Spectrum SpecularReflection::sample(BSDFRecord &bsdfRec, const Point2f &sample, Float &pdf) const {
        auto wi = bsdfRec.isct.shCoord.toLocal(bsdfRec.wo);
        wi.x = -wi.x;
        wi.y = -wi.y;
        CoordinateSystem wi_coord(wi);
        Vec3f offset = wi_coord.toWorld(cosineHemisphereSampling(sample));
        wi += offset * roughness;
        bsdfRec.wi = bsdfRec.isct.shCoord.toWorld(wi);
        bsdfRec.wi.normalize();
        pdf = 1.f;
        return {1.f};
    }

    Spectrum SpecularReflection::evaluate(BSDFRecord &bsdfRec, Float &pdf) const {
        pdf = this->pdf(bsdfRec);
        // TODO: ref-spec eval
        return {};
    }

    Float SpecularReflection::pdf(const BSDFRecord &bsdfRec) const {
        return .0f;
    }

    ///////////////////////////
    /// SpecularTransmission
    ///////////////////////////

    Spectrum SpecularTransmission::sample(BSDFRecord &bsdfRec, const Point2f &sample, Float &pdf) const {
        Vec3f ray_d = -bsdfRec.wo.normalized();
        Vec3f refracted;
        Vec3f reflected = reflect(ray_d, bsdfRec.isct.shN);

        float cosine;

        bool isRefracted;

        if (dot(ray_d, bsdfRec.isct.shN) > 0) {
            isRefracted = refract(ray_d, refracted, -bsdfRec.isct.shN, refract_rate);
            cosine = refract_rate * dot(ray_d, bsdfRec.isct.shN);
        } else {
            isRefracted = refract(ray_d, refracted, bsdfRec.isct.shN, 1.f / refract_rate);
            cosine = -dot(ray_d, bsdfRec.isct.shN);
        }

        if (isRefracted) {
            if (sample.x <= fresnel.evaluate(cosine).x) {
                isRefracted = false;
            }
        }
        
        bsdfRec.wi = isRefracted ? refracted : reflected;

        // float real_refract_rate = dot(bsdfRec.wo, bsdfRec.isct.shN) < 0.f ? refract_rate : 1.f / refract_rate;
        // auto wi = bsdfRec.isct.shCoord.toLocal(bsdfRec.wo);
        // float cosTheta = wi.z;
        // float sin2Theta = std::max(.0f, 1.f - cosTheta * cosTheta);
        // float sin2ThetaT = real_refract_rate * real_refract_rate * sin2Theta;
        // float cosTetaT = std::sqrt(1 - sin2ThetaT);
        // wi = real_refract_rate * -wi + (real_refract_rate * cosTheta - cosTetaT) * Vec3f{0, 0, 1.f};
        // bsdfRec.wi = bsdfRec.isct.shCoord.toWorld(wi);
        
        bsdfRec.wi.normalize();
        pdf = 1.f;
        return {1.f};
    }

    Spectrum SpecularTransmission::evaluate(BSDFRecord &bsdfRec, Float &pdf) const {
        pdf = this->pdf(bsdfRec);
        // TODO: trans-lambertian eval
        return {};
    }

    Float SpecularTransmission::pdf(const BSDFRecord &bsdfRec) const {
        return .0f;
    }

}
